package cn.edu.sdjzu.knoverse.config;

import dev.langchain4j.community.model.zhipu.ZhipuAiChatModel;
import dev.langchain4j.community.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.community.model.zhipu.ZhipuAiStreamingChatModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.internal.Utils;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.content.retriever.neo4j.Neo4jContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import dev.langchain4j.store.graph.neo4j.Neo4jGraph;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

@Configuration
public class LangChain4jConfig {

    @Bean
    public ChatLanguageModel getAiChatModel() {
        ChatLanguageModel chatLanguageModel = ZhipuAiChatModel.builder()
                .apiKey("7b2fcef5dd914fa28e5697778dad483f.M8JA0rCbiJV1N1M3")
//                .model("glm-4-air-250414")
                .model("glm-4-plus")
//                .model("glm-4-flash")
                .temperature(0.6)
                .maxToken(1024)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .logRequests(true)
                .logResponses(true)
                .build();
        return chatLanguageModel;
    }
    
    // 智谱AI流式聊天模型配置
    @Bean
    public StreamingChatLanguageModel getStreamingAiChatModel() {
        StreamingChatLanguageModel streamingChatLanguageModel = ZhipuAiStreamingChatModel.builder()
                .apiKey("7b2fcef5dd914fa28e5697778dad483f.M8JA0rCbiJV1N1M3")
//                .model("glm-4-flash")
                .model("glm-4-plus")
                .temperature(0.6)
                .maxToken(1024)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .logRequests(true)
                .logResponses(true)
                .build();
        return streamingChatLanguageModel;
    }

    // 嵌入模型配置
    @Bean
    public EmbeddingModel getEmbeddingModel() {
        ZhipuAiEmbeddingModel model = ZhipuAiEmbeddingModel.builder()
                .apiKey("7b2fcef5dd914fa28e5697778dad483f.M8JA0rCbiJV1N1M3")
                .logRequests(true)
                .logResponses(true)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .logRequests(true)
                .logResponses(true)
                .build();
        return model;
    }

    // PostgreSQL向量存储配置
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel) {
//        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
//                .host("localhost")  // Required: Host of the PostgreSQL instance
//                .port(5432)  // Required: Port of the PostgreSQL instance
//                .database("testdb")  // Required: Database name
//                .user("postgres")  // Required: Database user
//                .password("xiaoxin")  // Required: Database password
//                .table("knoverse")  // Required: Table name to store embeddings
//                .dimension(embeddingModel.dimension()) // Required: Dimension of embeddings
//                .build();
        MilvusEmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
                .host("192.168.6.129")                         // Host for Milvus instance
                .port(19530)                               // Port for Milvus instance
                .collectionName("knoverse_collection")      // Name of the collection
                .dimension(embeddingModel.dimension())                            // Dimension of vectors
                .indexType(IndexType.FLAT)                 // Index type
                .metricType(MetricType.COSINE)             // Metric type
//                .username("username")                      // Username for Milvus
//                .password("password")                      // Password for Milvus
                .consistencyLevel(ConsistencyLevelEnum.EVENTUALLY)  // Consistency level
                .autoFlushOnInsert(true)                   // Auto flush after insert
                .idFieldName("id")                         // ID field name
                .textFieldName("text")                     // Text field name
                .metadataFieldName("metadata")             // Metadata field name
                .vectorFieldName("vector")                 // Vector field name
                .build();                                  // Build the MilvusEmbeddingStore instance
        return embeddingStore;
    }
    // Neo4j驱动配置
    @Bean
    public Driver neo4jDriver() {
        return GraphDatabase.driver(
                "bolt://localhost:7687",  // Neo4j默认端口
                AuthTokens.basic("neo4j", "zf07130828")  // 用户名和密码
        );
    }

    // Neo4j图数据库配置
    @Bean
    public Neo4jGraph neo4jGraph(Driver driver) {
        return Neo4jGraph.builder()
                .driver(driver)
                .build();
    }

    // Neo4j内容检索器配置
    @Bean
    public Neo4jContentRetriever neo4jContentRetriever(Neo4jGraph graph, ChatLanguageModel chatLanguageModel) {
//        PromptTemplate promptTemplate = PromptTemplate.from("Based on the Neo4j graph schema below, write a Cypher query that would answer the user's question." +
//                "Note: You can use the APOC extension for fuzzy matching of node names. Minimize the number of query conditions to match more nodes. For users’ colloquial expressions, queries must convert the terms accordingly and apply fuzzy matching. For example, when a user asks 'What are the knowledge points in the Soft Eng course?', the full term is 'Software Engineering', so the query should perform a fuzzy search using 'Software Engineering'.\n" +
//                "Output only pure Cypher statements without any decoration. Do not use Markdown syntax, and it is prohibited to output non - query statements. If there is no corresponding statement, output a Cypher statement that returns an empty result.\n" +
//                "Schema: {{schema}}\n\nQuestion: {{question}}\nCypher query:\n");
//        PromptTemplate promptTemplate = PromptTemplate.from("根据下面的 Neo4j 图数据库模式，编写一个 Cypher 查询语句来回答用户的问题。" +
//                "注意：你可以使用 APOC 扩展来模糊匹配节点名称。尽量减少查询条件的数量，以匹配更多的节点。对于用户的口语表达，需要转换术语并进行模糊匹配。例如，当用户问“软工课程有哪些知识点？”时，完整术语是“软件工程”，因此查询应使用“软件工程”进行模糊搜索。\n" +
//                "仅输出纯粹的 Cypher 查询语句，不允许有任何修饰语句。不使用 Markdown 语法，禁止输出非查询类语句。如果没有对应的查询语句，请输出一个返回空结果的 Cypher 查询语句。\n" +
//                "Schema: {{schema}}\n\nQuestion: {{question}}\nCypher query:\n");
        PromptTemplate promptTemplate = PromptTemplate.from("你的任务是根据以下 Neo4j 图数据库模式，编写一个 Cypher 查询语句以回答用户的问题。\n\n" +
                "在编写查询时，请遵循以下规则：\n" +
                "1. 可以使用 APOC 扩展来实现节点名称的模糊匹配。\n" +
                "2. 尽量减少查询条件，以匹配更多相关节点。\n" +
                "3. 对于用户的口语表达，应转换为标准术语并进行模糊匹配。例如，用户问“软工课程有哪些知识点？”，应识别为“软件工程”并进行模糊搜索。\n" +
                "4. 识别用户问题中的缩写、别名或简写，并自动映射为图中的标准术语。\n" +
                "5. 如果无法生成合理的查询语句，请返回一个有效的空结果查询，例如：`RETURN null LIMIT 0`\n" +
                "6. 严禁输出任何非查询语句内容（如注释、说明、解释等），不得使用 Markdown 或格式修饰。\n\n" +
                "Schema: {{schema}}\n\nQuestion: {{question}}\nCypher query:\n");

        return Neo4jContentRetriever.builder()
                .graph(graph)
                .chatLanguageModel(chatLanguageModel)
                .promptTemplate(promptTemplate)
                .build();
    }

    // 向量存储内容检索器配置
    @Bean
    public ContentRetriever vectorContentRetriever(EmbeddingStore<TextSegment> embeddingStore, 
                                                 EmbeddingModel embeddingModel) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3)
                .minScore(0.6)
                .build();
    }
}
